import json
import time

import gradio as gr
import numpy as np
from ocr_process.cls.clas_pipeline import ClasPipeline
from ocr_process.operators import ClasResizeImage
from ocr_process.utils.img_utils import np2base64
from openai import OpenAI

from taskinfo.uv_pipeline import UvPipeline

uv_predict = UvPipeline(model_path=r"D:\ocr\onnx\deploy\uvdoc\inference.onnx", use_openvino=False)
clas_predict = ClasPipeline(model_path=r"D:\ocr\train_data\multi_class\jinshi_demo\inference.onnx", use_openvino=False)
resize_func = ClasResizeImage(960)


# resize_func = ClasResizeImage(960)
def predict(image: np.ndarray, prompt: str):
    # print(f"prompt={prompt}")
    result = uv_predict(image)
    if result['output_imgs']:
        image = result['output_imgs'][0]

    # img = Image.fromarray(image_data)
    # r_channel = img.split()[0]
    # rgb_image = r_channel.convert('RGB')
    # image_data = np.array(rgb_image)
    print(f"original shap={image.shape}")
    start_time = time.perf_counter()
    if image.shape[0] > 960 or image.shape[1] > 960:
        image_data = resize_func({"image": image})["image"]
        # print(f"resize shap={image_data.shape}")
    text_tesult = inference_with_api(np2base64(image), prompt=prompt)
    print(f"耗时：{time.perf_counter() - start_time}")
    return text_tesult




def inference_with_api(base64_image, prompt, sys_prompt="你是一个文档信息提取助手，提取文档中信息来回答问题。",
                       model_id="qwen2.5-vl-7b-instruct"):
    """
    sk-bfc7ee0784464456b6d3de680afdd4dd
    https://dashscope.aliyuncs.com/compatible-mode/v1
    qwen2.5-vl-7b-instruct
    """
    client = OpenAI(
        api_key="sk-bfc7ee0784464456b6d3de680afdd4dd",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    messages = [
        {"role": "system", "content": sys_prompt},
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                },
                # {"type": "text", "text": prompt},
                {"type": "text", "text": f"文档信息提取,提取内容只能包含'{prompt}',请你记住你的输出是一个json格式的字符串"},
                # {"type": "text", "text": f"文档信息提取,提取内容只能包含'{prompt}'，时间转换成格式为yyyy-MM-dd HH:mm:ss,请你记住你的输出是一个json格式的字符串"},
            ],
        }
    ]
    completion = client.chat.completions.create(
        model=model_id,
        messages=messages,
        temperature=0.7,

    )
    # print(completion)
    content: str = completion.choices[0].message.content
    try:
        bot_response = content[content.find("{"):]
        bot_response = bot_response[: bot_response.find("}") + 1]
        return bot_response
    except:
        pass
    return content


prompt_dict = {
    "金隅水泥": "订单编号,称重时间,打印时间,材料品种,供货单位,车牌号码,毛重,皮重,净重",
    "内蒙古杭萧盛": "车号,毛重,皮重,净重,过毛时间,过皮时间,扣重",
    "内蒙古汇能": "车牌号,毛重,皮重,净重,过重时间,接单时间",
    "四川发展": "品名,承运车号,毛重,车重,净重,实收吨重",
}


def gradio_interface(image):
    image_np = np.array(image)
    input_data = {"image": image_np, "original_image": image_np.copy()}
    # 图像分类预测
    clas_res = clas_predict(input_data)
    print(f"res = {clas_res}")
    class_id = clas_res["label_names"][0]
    scores = clas_res["scores"][0]
    if scores < 0.78:
        return "目前不支持该类型图片识别，请重新上传【金隅水泥, 内蒙古杭萧盛, 内蒙古汇能, 四川发展】"

    json_str = predict(image, prompt_dict[class_id])
    res = {"type": class_id, "result": json.loads(json_str)}
    return json.dumps(res, ensure_ascii=False, indent=4)


iface = gr.Interface(
    fn=gradio_interface,
    inputs=["image"],
    outputs=['text'],
    title="磅单识别演示【金隅水泥, 内蒙古杭萧盛, 内蒙古汇能, 四川发展】",
    description="上传一张磅单，识别信息"
)

if __name__ == '__main__':
    iface.launch(server_name="0.0.0.0")
